#!/usr/bin/python
# -*- coding:utf-8 -*-
"""
备份tasks表数据到备份数据库，备份完成后会删除原数据
备份 3天前的数据， 可用性任务 （taskType in [http, ping]）备份到 available_tasks 表
非可用性任务 （taskType not in [http, ping]）备份 tasks 表
删除可用性任务75天前，非可用性任务210天前的数据
"""

import time
import datetime
from bson.objectid import ObjectId
from bson.py3compat import abc
from bson.raw_bson import RawBSONDocument

from pymongo.message import _BulkWriteContext, _EncryptedBulkWriteContext
from pymongo.bulk import _Bulk
from pymongo import MongoClient
from pymongo import common, message
from pymongo.results import InsertManyResult
from pymongo.database import Database
from pymongo.collection import Collection

from conf.settings import (
    CLOUDMONITOR_MONGODB,
    CLOUDMONITOR_DB,
    CLOUDMONITOR_MONGODB_BK,
    CLOUDMONITOR_DB_BK
)
from utility.wx_notify import notify
from utility.log import Logger

logger = Logger(filename='tasks_collection_bk.log').get_logger()


class _CustomBulkWriteContext(_BulkWriteContext):
    """重写一下check_keys方法， 使批量插入可以接受key包含 . 的文档 """

    @property
    def check_keys(self):
        """Should we check keys for this operation type?"""
        return False


class _CustomEncryptedBulkWriteContext(_EncryptedBulkWriteContext):

    @property
    def check_keys(self):
        """Should we check keys for this operation type?"""
        return False


class _CustomBulk(_Bulk):

    @property
    def bulk_ctx_class(self):
        encrypter = self.collection.database.client._encrypter
        if encrypter and not encrypter._bypass_auto_encryption:
            return _CustomEncryptedBulkWriteContext
        else:
            return _CustomBulkWriteContext


class CustomCollection(Collection):

    def insert_many(self, documents, ordered=True, bypass_document_validation=False, session=None):

        if not isinstance(documents, abc.Iterable) or not documents:
            raise TypeError("documents must be a non-empty list")
        inserted_ids = []

        def gen():
            """A generator that validates documents and handles _ids."""
            for document in documents:
                common.validate_is_document_type("document", document)
                if not isinstance(document, RawBSONDocument):
                    if "_id" not in document:
                        document["_id"] = ObjectId()
                    inserted_ids.append(document["_id"])
                yield (message._INSERT, document)

        write_concern = self._write_concern_for(session)
        blk = _CustomBulk(self, ordered, bypass_document_validation)
        blk.ops = [doc for doc in gen()]
        blk.execute(write_concern, session=session)
        return InsertManyResult(inserted_ids, write_concern.acknowledged)


class CustomDatabase(Database):

    def __getitem__(self, name):
        return CustomCollection(self, name)


class CustomMongoClient(MongoClient):

    def __getitem__(self, name):
        return CustomDatabase(self, name)


class BackupData(object):

    def __init__(self):
        self.db = CustomMongoClient(CLOUDMONITOR_MONGODB)[CLOUDMONITOR_DB]
        self.bk_db = CustomMongoClient(CLOUDMONITOR_MONGODB_BK)[CLOUDMONITOR_DB_BK]
        # 数据存储时间为utc时间
        self.now = datetime.datetime.now() - datetime.timedelta(hours=8)
        self.available_task_time = self.now - datetime.timedelta(days=75)
        self.other_task_time = self.now - datetime.timedelta(days=210)
        self.three_days_ago_time = self.now - datetime.timedelta(days=3)
        self.seven_days_ago_time = self.now - datetime.timedelta(days=7)
        self.available_task_id = ObjectId('%x' % int(self.available_task_time.timestamp()) + '0'*16)
        self.other_task_id = ObjectId('%x' % int(self.other_task_time.timestamp()) + '0'*16)
        self.three_days_ago_id = ObjectId('%x' % int(self.three_days_ago_time.timestamp()) + '0'*16)
        self.seven_days_ago_id = ObjectId('%x' % int(self.seven_days_ago_time.timestamp()) + '0'*16)

    def get_last_tail_task_id(self, task_type="available"):
        """查询上一次备份的最后一条数据"""
        collection = "tasks" if task_type != "available" else "available_tasks"
        cursor = list(self.bk_db[collection].find({}, {"_id": 1}).sort("_id", -1).limit(1))
        if cursor:
            return cursor[0]["_id"]
        return None

    def get_db1_oldest_id(self):
        """查询业务数据库中保留着的最早一条数据"""
        cursor = list(self.db.tasks.find({}, {"_id": 1}).sort("_id", 1).limit(1))
        return cursor[0]["_id"]

    def _backup_data(self, collection, task_list, save_count):
        result = self.bk_db[collection].insert_many(task_list)
        inserted_ids = result.inserted_ids
        save_count += len(inserted_ids)
        return save_count

    def backup_data(self, cursor, task_type="available"):
        """备份和删除数据"""
        save_count = 0
        total_count = 0
        collection = "tasks" if task_type != "available" else "available_tasks"
        task_list = []
        for task in cursor:
            if len(task_list) >= 100:
                save_count = self._backup_data(collection, task_list, save_count)
                task_list = []
            task_list.append(task)
            total_count += 1
        if task_list:
            save_count = self._backup_data(collection, task_list, save_count)
        return total_count, save_count

    def _update_not_completed_task_bk(self, collection, task_list):
        """先在备份数据库中删除之前的数据， 然后插入tasks表中新的数据"""
        delete_ids = [t["_id"] for t in task_list]
        task_list = list(self.db["tasks"].find({"_id": {"$in": delete_ids}}))
        if not task_list:
            return
        self.bk_db[collection].delete_many({"_id": {"$in": delete_ids}})
        self.bk_db[collection].insert_many(task_list)

    def update_not_completed_task_bk(self, task_type="available"):
        """处理上次备份数据中不是 completed 或者 stopped 任务数据, 取一个星期的数据"""
        db1_oldest_id = self.seven_days_ago_id
        collection = "tasks" if task_type != "available" else "available_tasks"
        cursor = self.bk_db[collection].find(
            {"_id": {"$gte": db1_oldest_id}, "status": {"$nin": ["completed", "stopped"]}}, {"_id": 1})
        task_list = []
        for task in cursor:
            if len(task_list) >= 100:
                self._update_not_completed_task_bk(collection, task_list)
                task_list = []
            task_list.append(task)
        if task_list:
            self._update_not_completed_task_bk(collection, task_list)

    def delete_old_data(self):
        """删除老的数据"""
        query1 = {"_id": {"$lt": self.available_task_id}, "taskType": {"$in": ["ping", "http"]}}
        query2 = {"_id": {"$lt": self.other_task_id}, "taskType": {"$nin": ["ping", "http"]}}
        r1 = self.db["tasks"].delete_many(query1)
        r2 = self.db["tasks"].delete_many(query2)
        return r1.deleted_count + r2.deleted_count

    def run(self):
        logger.info("开始数据库备份......")

        # 先处理上次备份中不是完成状态的任务
        logger.info("更新之前没有完成的任务状态......")
        for kind in ["other", "available"]:
            self.update_not_completed_task_bk(kind)
        logger.info("备份新的数据......")
        start_time = int(time.time())
        available_task_last_id = self.get_last_tail_task_id()
        other_task_last_id = self.get_last_tail_task_id(task_type="other")
        available_task_query = {
            "taskType": {"$in": ["http", "ping"]},
            "_id": {"$lte": ObjectId(self.three_days_ago_id)},
        }
        other_task_query = {
            "taskType": {"$nin": ["http", "ping"]},
            "_id": {"$lte": ObjectId(self.three_days_ago_id)},
        }
        if available_task_last_id:
            available_task_query["_id"].update({"$gt": available_task_last_id})
        if other_task_last_id:
            other_task_query["_id"].update({"$gt": other_task_last_id})
        available_task_cursor = self.db["tasks"].find(available_task_query)
        other_task_cursor = self.db["tasks"].find(other_task_query)
        available_task_count, available_task_save_count = self.backup_data(
            available_task_cursor)
        other_task_count, other_task_save_count = self.backup_data(
            other_task_cursor, "other")

        logger.info("删除旧的数据......")
        deleted_count = self.delete_old_data()

        log_message = f"数据备份结束, 用时{int(time.time() - start_time)}秒, " \
                      f"需要备份可用性任务{available_task_count}条， 非可用性任务{other_task_count}条,  " \
                      f"实际备份完成可用性任务{available_task_save_count}条， 非可用性任务{other_task_save_count}条, " \
                      f"共删除历史数据{deleted_count}条"
        logger.info(log_message)
        notify(log_message)


if __name__ == "__main__":
    BackupData().run()
